# example/01_esay/04_relu与二次拟合.py

import sys
sys.path.append('../..')  # 父目录的父目录
from ourdl.core import Varrible
from ourdl.ops import Mul, Add
from ourdl.ops.loss import ValueLoss
from ourdl.ops import LeakyRelu as Relu
import matplotlib.pyplot as plt
import numpy as np
import random

# 1.1 线性变换一
x = Varrible()
w_11 = Varrible()
w_12 = Varrible()
mul_11 = Mul([x, w_11])
mul_12 = Mul([x, w_12])
b_11 = Varrible()
b_12 = Varrible()
add_11 = Add([mul_11, b_11])
add_12 = Add([mul_12, b_12])
# 1.2 激活函数 --> 非线性变换
relu_11 = Relu([add_11])
relu_12 = Relu([add_12])
# 1.3 线性变换二
w_21 = Varrible()
w_22 = Varrible()
mul_21 = Mul([relu_11, w_21])
mul_22 = Mul([relu_12, w_22])
b_21 = Varrible()
add_21 = Add([mul_21, mul_22, b_21])
# 1.4 损失函数
label = Varrible()
loss = ValueLoss([label, add_21])

# 2 参数初始化
params = [w_11, w_12, b_11, b_12, w_21, w_22, b_21]
for param in params:
    param.set_value(random.uniform(-1, 1))
print([param.value for param in params])

# 3 生成数据
data_x = [random.uniform(0, 2) for i in range(1500)]  # 似乎实数比离散的[0, 1, 2]要好
data_label = [x * x for x in data_x]

# 4 开始训练
# 4.1 训练，同时绘制动画
losses = []
for i in range(len(data_x)):
    x.set_value(data_x[i])
    label.set_value(data_label[i])
    loss.forward()
    for param in params:
        param.get_grad()
        param.update(lr=0.01)
    if i % 100 == 0:
        print(f'[{i}]:loss={loss.value},', [param.value for param in params])
    losses.append(loss.value)
    loss.clear()

# 5 画出训练过程中loss的变化曲线
show_x = [i for i in range(len(losses))]
show_y = [_ for _ in losses]
plt.plot(show_x, show_y)
plt.show()